#include <iostream>

namespace me
{
    enum class color
    {
        RED,
        BLACK
    };

    template <typename K, typename V>
    struct rbtreeNode
    {
        std::pair<K, V> _kv;
        rbtreeNode* _left;
        rbtreeNode* _right;
        rbtreeNode* _parent;
        color _col;

        rbtreeNode(const std::pair<K, V>& kv)
            :_kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr)
        {}
    };

    template <typename K, typename V>
    class rbtree
    {
        using node = rbtreeNode<K, V>;
        void destroy(node* root)
        {
            if(root == nullptr) return;
            destroy(root->_left);
            destroy(root->_right);
            delete root;
            root = nullptr;
        }

        node* copy(const node* other)
        {
            if(other == nullptr) return nullptr;
            node* newnode = new node(other->_kv);
            newnode->_col = other->_col;

            newnode->_left = copy(other->_left);
            newnode->_right = copy(other->_right);

            if(newnode->_left) newnode->_left->_parent = newnode;
            if(newnode->_right) newnode->_right->_parent = newnode;

            return newnode;
        }
    public:
        rbtree():_root(nullptr) {}
        ~rbtree() { destroy(_root); }
        rbtree(const rbtree& other) { _root = copy(other._root); }
        rbtree& operator=(rbtree other)
        {
            std::swap(_root, other._root);
            return *this;
        }

        bool insert(const std::pair<K, V>& kv)
        {
            // 如果是空树，直接将新节点赋值给root
            if(_root == nullptr) {
                _root = new node(kv);
                _root->_col = color::BLACK; // 根节点一定是黑色的
                return true;
            }

            // 1.根据二叉树的插入逻辑，进行左右比较，插入到合适位置
            node* parent = nullptr;
            node* cur = _root;
            while(cur) {
                if(kv.first < cur->_kv.first) {
                    parent = cur;
                    cur = cur->_left;
                }else if(kv.first > cur->_kv.first) {
                    parent = cur;
                    cur = cur->_right;
                }else {
                    return false; // 不允许冗余
                }
            }

            cur = new node(kv);
            cur->_col = color::RED; // 新插入的节点除了根节点外必须是红色的
            if(kv.first < parent->_kv.first) {
                parent->_left = cur;
            }else {
                parent->_right = cur;
            }
            cur->_parent = parent;

            // 2.如果出现连续的红色，根据分类进行更新
            // 当g是根时，更新完成后不用再更新了，此时cur = root，parent = nullptr，此时直接退出，并且将根设置为黑色
            while(parent && parent->_col == color::RED) {
                node* grand = parent->_parent;

                if(grand->_left == parent) {
                    //       g
                    //    p     u
                    node* uncle = grand->_right;

                    // 当uncle存在且为红色，变色
                    // g变红，p、u变黑
                    if(uncle && uncle->_col == color::RED) {
                        parent->_col = color::BLACK;
                        uncle->_col = color::BLACK;
                        grand->_col = color::RED;

                        // 更新cur、parent，继续向上更新
                        cur = grand;
                        parent = cur->_parent;
                    }else {
                        // 当uncle不存在，或者存在为黑色
                        if(cur == parent->_left) {
                            //       g
                            //    p     u
                            // c
                            // 当一边高时，进行单旋+变色 :p变黑，g变红
                            rotateR(grand);
                            parent->_col = color::BLACK;
                            grand->_col = color::RED;
                        }else {
                            //       g
                            //    p     u
                            //      c
                            // 当出现右左不均时，进行双旋+变色：c变黑，g变红
                            rotateL(parent);
                            rotateR(grand);
                            cur->_col = color::BLACK;
                            grand->_col = color::RED;
                        }
                        break; // 旋转+变色之后不需要继续更新了
                    }
                }else {
                    //       g
                    //    u     p
                    node* uncle = grand->_left;

                    // uncle存在且为红，变色：u、p变黑，g变红
                    if(uncle && uncle->_col == color::RED) {
                        parent->_col = color::BLACK;
                        uncle->_col = color::BLACK;
                        grand->_col = color::RED;

                        // 更新cur、parent，继续向上更新
                        cur = grand;
                        parent = cur->_parent;
                    }else {
                        // uncle不存在/uncle存在为黑色，根据插入位置的不同选择单旋或者双旋，并且进行变色
                        if(cur == parent->_right) {
                            //       g
                            //    u     p
                            //             c
                            rotateL(grand);
                            // 单旋：p变黑，g变红
                            parent->_col = color::BLACK;
                            grand->_col = color::RED;
                        }else {
                            //       g
                            //    u     p
                            //        c
                            rotateR(parent);
                            rotateL(grand);
                            // 双旋：c变黑，g变红
                            cur->_col = color::BLACK;
                            grand->_col = color::RED;
                        }
                        break;
                    }
                }
            }
            _root->_col = color::BLACK; // 根必须是黑色的
            return true;
        }

        node* find(const K& key)
        {
            node* cur = _root;
            while(cur) {
                if(key < cur->_kv.first) {
                    cur = cur->_left;
                }else if(key > cur->_kv.first) {
                    cur = cur->_right;
                }else {
                    return cur;
                }
            }
            return nullptr;
        }
    private:
        // 右单旋
        void rotateR(node* prev)
        {
            //       pprev
            //       prev
            //   cur
            //     curR
            node* cur = prev->_left;
            node* curR = cur->_right;
            node* pprev = prev->_parent;

            //       pprev
            //       cur
            //            prev
            //          curR

            // 1.cur的右子树，变为prev的左子树
            if(curR) curR->_parent = prev; 
            prev->_left = curR;

            // 2. prev变为cur的右子树
            cur->_right = prev;
            prev->_parent = cur;

            // 3.更新cur的指向
            if(prev == _root) {
                _root = cur;
                cur->_parent = nullptr;
            }else {// 如果prev原先不是根，则判断cur应该连接到pprev的左子树还是右子树
                if(prev == pprev->_left) {
                    pprev->_left = cur;
                }else{
                    pprev->_right = cur;
                }
                cur->_parent = pprev;
            }
        }

        // 左单旋
        void rotateL(node* prev)
        {
            //   pprev
            //   prev
            //       cur
            //      curL
            node* cur = prev->_right;
            node* curL = cur->_left;
            node* pprev = prev->_parent;

            //      pprev
            //      cur
            //  prev
            //    curL

            // 1.cur的左子树，成为prev的右子树
            if(curL) curL->_parent = prev;
            prev->_right = curL;

            // 2.prev变为cur的左子树
            cur->_left = prev;
            prev->_parent = cur;

            // 3.更新cur的指向
            if(prev == _root) {
                _root = cur;
                cur->_parent = nullptr;
            }else {
                if(prev == pprev->_left) {
                    pprev->_left = cur;
                }else {
                    pprev->_right = cur;
                }
                cur->_parent = pprev;
            }
        }
    private:
        node* _root;
    };
}